import os
import numpy as np
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.cluster import KMeans

# 读取文本文件
def read_file(file_path):
    with open(file_path, 'r', encoding='utf-8') as f:
        data = f.readlines()[:500]
        return data

# 加载数据
def load_data():
    data_path = "test1.txt"  # 文件路径
    statements = read_file(data_path)

    drug_statements = {}
    # 文本分析或正则表达式筛选涉及毒品的事实，并按不同的毒品类别存储
    for statement in statements:
        if "冰毒" in statement:
            if "冰毒" in drug_statements:
                drug_statements["冰毒"].append(statement)
            else:
                drug_statements["冰毒"] = [statement]
        elif "海洛因" in statement:
            if "海洛因" in drug_statements:
                drug_statements["海洛因"].append(statement)
            else:
                drug_statements["海洛因"] = [statement]
        elif "k粉" in statement:
            if "k粉" in drug_statements:
                drug_statements["k粉"].append(statement)
            else:
                drug_statements["海洛因"] = [statement]
        elif "神仙水" in statement:
            if "神仙水" in drug_statements:
                drug_statements["神仙水"].append(statement)
            else:
                drug_statements["神仙水"] = [statement]
        elif "甲基苯丙胺" in statement:
            if "甲基苯丙胺" in drug_statements:
                drug_statements["甲基苯丙胺"].append(statement)
            else:
                drug_statements["甲基苯丙胺"] = [statement]
        elif "开心果" in statement:
            if "开心果" in drug_statements:
                drug_statements["开心果"].append(statement)
            else:
                drug_statements["开心果"] = [statement]
        elif "咖啡因" in statement:
            if "咖啡因" in drug_statements:
                drug_statements["咖啡因"].append(statement)
            else:
                drug_statements["咖啡因"] = [statement]
        elif "可卡因" in statement:
            if "可卡因" in drug_statements:
                drug_statements["可卡因"].append(statement)
            else:
                drug_statements["可卡因"] = [statement]
        elif "大麻" in statement:
            if "大麻" in drug_statements:
                drug_statements["大麻"].append(statement)
            else:
                drug_statements["大麻"] = [statement]
        elif "麻古" in statement:
            if "麻古" in drug_statements:
                drug_statements["麻古"].append(statement)
            else:
                drug_statements["麻古"] = [statement]
        elif "吗啡" in statement:
            if "吗啡" in drug_statements:
                drug_statements["吗啡"].append(statement)
            else:
                drug_statements["吗啡"] = [statement]
        elif "摇头丸" in statement:
            if "摇头丸" in drug_statements:
                drug_statements["摇头丸"].append(statement)
            else:
                drug_statements["摇头丸"] = [statement]

    return drug_statements


# 文本向量化
def vectorize(texts):
    vectorizer = TfidfVectorizer(stop_words=None, max_features=2000)
    vectors = vectorizer.fit_transform(texts).toarray()
    return vectors


# 对每个类别进行K均值聚类
def cluster_each_type(drug_statement_pairs):
    drug_type = drug_statement_pairs[0]
    drug_statements = drug_statement_pairs[1]

    if len(drug_statements) < 2:
        return {}

    data = vectorize(drug_statements)
    labels = KMeans(n_clusters=2).fit_predict(data)

    clusters = {}
    for i, statement in enumerate(drug_statements):
        if labels[i] in clusters:
            clusters[labels[i]].append(statement.strip())
        else:
            clusters[labels[i]] = [statement.strip()]

    return {drug_type: clusters}


# 主程序
def main():
    drug_statements = load_data()
    clusters = {}
    for drug_type, statements in drug_statements.items():
        clusters.update(cluster_each_type((drug_type, statements)))

    # 打印结果
    for drug_type, clusters_dict in clusters.items():
        print("毒品类型：{}".format(drug_type))
        for i, (label, cluster) in enumerate(clusters_dict.items()):
            print("聚类{}:".format(i))
            print("\n".join(cluster))
        print("-" * 50)


if __name__ == '__main__':
    main()
